# -*- coding: utf-8 -*-
"""
Created on Fri Jun 27 10:59:17 2025

@author: Administrator
"""

import arcpy
import socket
import struct


def grid_server(input_grid):
    input_raster = arcpy.Raster(input_grid)
    # data_type = ["U1", "U2", "U4", "U8", "S8", "U16",
    # "S16", "U32", "S32", "F32", "F64"]
    # 获取栅格的属性，根据数据的第一个波段，如左下角坐标、像元大小、行列数、投影信息等
    desc = arcpy.Describe(input_grid)  # 描述信息
    desc1 = arcpy.Describe(input_grid + '/Band_1')    # 描述信息1
    x_origin = desc.extent.XMin     # 最小X
    y_origin = desc.extent.YMin     # 最小Y
    cell_size = input_raster.meanCellWidth  # 像元大小
    rows, cols = desc1.height, desc1.width  # 行列数
    pixel_type = desc1.pixelType  # 像素类型和深度
    band_count = desc.bandCount  # 波段数
    noDataValue = desc1.noDataValue  # 空值数
    dx = rows*rows*band_count/10**9       # 预估内存大小
    # 像素类型和像素深度字典,获取后面需要的类型操作符
    pixel_typea = {
        "U8": ["uint8", "8_BIT_UNSIGNED", 'B'],
        "S8": ["int8", "8_BIT_SIGNED", 'b'],
        "U16": ["uint16", "16_BIT_UNSIGNED", 'H'],
        "S16": ["int16", "16_BIT_SIGNED", 'h'],
        "U32": ["uint32", "32_BIT_UNSIGNED", 'I'],
        "S32": ["int32", "32_BIT_SIGNED", 'i'],
        "F32": ["float32", "32_BIT_FLOAT", 'f'],
        "F64": ["float64", "64_BIT", 'd']}
    if pixel_type in pixel_typea:
        d_tpye1 = pixel_typea[pixel_type][0]   # uint8
        d_tpye2 = pixel_typea[pixel_type][1]   # 8_BIT_UNSIGNED
        type_codes = pixel_typea[pixel_type][2]   # B
    print(f"\n栅格数据的像素位深是：{d_tpye1}\n编码类型是：{d_tpye2}\nstruct类型码是：{type_codes}")

    # 获取投影的 WKT 字符串
    spatial_ref = input_raster.spatialReference  # 获取投影信息
    sr_wkt = spatial_ref.exportToString()  # 获取投影信息的字符串表达
    sr_wkt_bytes = sr_wkt.encode('utf-8')  # 将投影信息转换为字节
    sr_wkt_len = len(sr_wkt)         # 获取字符串的长度：这里等于530

    # 输出信息供用户选择
    print(
        f'new原始栅格数据的左下角坐标为: {x_origin} , {y_origin}',
        f'原始栅格数据的列数和行数分别为: 列：{cols} , 行：{rows}',
        f'原始栅格数据的波段数为: {band_count}',
        f'原始栅格数据的像素位深及类型为: {pixel_type}',
        f'原始栅格数据的None空值为: {noDataValue}',
        '=========================================================',
        '下面将根据数据情况进行分块处理，每块大小根据实际内存决定',
        '具体计算方法：行数*列数*3（波段数）除以10的九次方，',
        f'如：{cols}*{rows}*3/10^9≈{dx} GB大小',
        '下面的分块大小不要超过剩余内存的50%',
        sep='\n'
        )

    # 判断用户输入行列参数和分块序号
    while True:
        ncols1 = int(input("请输入分块列大小: "))
        nrows1 = int(input("请输入分块行大小: "))
        blocknum_x, blocknum_y = map(int, input("请输入两个分块编号（两个数字，用空格分隔）：").split())
        print(f"\n你输入的分块列和行大小是：{ncols1},{nrows1}")
        print(f"你输入的分块序号为：{blocknum_x},{blocknum_y}")
        confirm = input("请检查上面输入是否正确？(y/n)：")
        if confirm == 'y':
            break  # 确认正确，退出循环
        elif confirm != 'n':
            print("⚠️ 请输入 'y' 或 'n'！")
            continue  # 忽略无效输入，重新确认
    print("\n✅ 所有输入已确认，程序继续执行......")

    lowerLeft = arcpy.Point(x_origin + blocknum_x * ncols1 * cell_size,
                            y_origin + blocknum_y * nrows1 * cell_size)

    # 转换为数组，增加数据类型设置和增加空值设置，要不要astype(d_tpye1)似乎不影响
    # 不同波段数栅格得到的数组维数不一样！！！
    raster_array = arcpy.RasterToNumPyArray(
        input_raster,
        lowerLeft,
        ncols=ncols1,
        nrows=nrows1,
        nodata_to_value=noDataValue).astype(d_tpye1)
    # 创建一个 TCP/IP 套接字
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    # 绑定地址和端口
    server_address = ('127.0.0.1', 12340)
    server_socket.bind(server_address)
    # 监听连接
    server_socket.listen(1)
    print('▲▲▲▲等待客户端连接▲▲▲▲......')

    # 接受客户端连接
    connection, client_address = server_socket.accept()
    print(f'连接来自: {client_address}')

    # 1-发送数组形状（修改：单波段栅格的shape是二维的，需要调整为三维表示）
    if band_count == 1:
        rows_block, cols_block = raster_array.shape
        shape = (band_count, rows_block, cols_block)  # 明确指定波段数为no
    else:
        shape = raster_array.shape
    connection.sendall(struct.pack('!III', *shape))

    # ++++++++++++++++++++++++++++++++++++++++++++++++++
    # 2-先将d_tpye2字符串编码为字节，并发送其长度和内容
    d_tpye2_bytes = d_tpye2.encode('utf-8')
    connection.sendall(struct.pack('!I', len(d_tpye2_bytes)))  # 发送长度
    connection.sendall(d_tpye2_bytes)                          # 发送内容

    # 3-发送 noDataValue（需根据其数据类型选择合适的 struct 格式，假设为整数）
    # 例如：若 noDataValue 是整数，用 'i' 或 'I'；若是浮点数，用 'f' 或 'd'
    # 注意：需确保发送端和接收端的数据类型一致，此处假设为 int32
    connection.sendall(struct.pack('!i', noDataValue))
    # ++++++++++++++++++++++++++++++++++++++++++++++++++

    # bit_depth = int(pixel_type[1:])  # 从pixel_type提取数字，如'U16'->16
    # connection.sendall(struct.pack('!I', bit_depth))  # 发送4字节整数

    # 4-1发送数组长度数据
    raster_array_bytes = raster_array.tobytes()  # 栅格数据字节化
    connection.sendall(struct.pack('!Q', len(raster_array_bytes)))
    # 4-2发送数组数据
    connection.sendall(raster_array_bytes)

    # 5-1发送投影信息长度
    connection.sendall(struct.pack('!I', sr_wkt_len))
    # 5-2发送投影信息内容
    connection.sendall(sr_wkt_bytes)

    # 6-发送属性，pack为打包成二进制数据
    attributes = [x_origin + blocknum_x * ncols1 * cell_size,
                  y_origin + blocknum_y * nrows1 * cell_size,
                  cell_size]
    attributes_bytes = struct.pack('!ddd', *attributes)
    connection.sendall(attributes_bytes)

    # 关闭连接
    connection.close()
    server_socket.close()
    print('连接关闭！')
    return
